[WIP] LM Workload #860
[WIP] LM Workload #860
Conversation
Dev -> main
Updates to LM PR
- Added `limit_tf_threads` parameter to `pytorch_init` to control TensorFlow threading based on workload type. Dataloader was going OOM otherwise. - Updated input pipeline to support "None" for weights (for memory). - Modified Transformer model's `forward` method to optionally return loss during training. Should be better to fuse the loss later. - Adjusted torch LM workload configuration for model dimensions and parameters to match jax. - Updated transformers version in `pyproject.toml`, older version seems unavailable.
…ency into lm_workload
…ency into lm_workload
… pytorch calls detatch
|
Adding some TODOs:
|
…JAX and PyTorch, also unify initialization to be the same in both
…uding learned scaling factor
There was a problem hiding this comment.
Second round of small requested changes.
Perhaps something we should discuss, we need a more descriptive name for the workload. E.g. fineweb_edu_lm. What do you all think? @Niccolo-Ajroldi @rka97
| @@ -0,0 +1,397 @@ | |||
| """ | |||
| Originally based on code from the NanoDO repository under the Apache 2.0 license: | |||
There was a problem hiding this comment.
Can we rename this file to models.py to be consistent with the pattern in the other workload definitions.
| @@ -0,0 +1,344 @@ | |||
| """ | |||
| Originally based on the plainLM codebase: | |||
There was a problem hiding this comment.
Can we rename this file to models.py to be consistent with the other workload definitions
| 'workload_path': 'librispeech_deepspeech/librispeech', | ||
| 'workload_class_name': 'LibriSpeechDeepSpeechNormAndSpecAugWorkload', | ||
| }, | ||
| 'lm': {'workload_path': 'lm/lm', 'workload_class_name': 'LmWorkload'}, |
There was a problem hiding this comment.
Now that we have all the important implementation details figured out should we pick a more descriptive name for the workload? I am thinking perhaps 'fineweb_edu_lm'?
| elif workload_name == 'mnist': | ||
| return 16 | ||
| elif workload_name == 'lm': | ||
| return 4 |
There was a problem hiding this comment.
This should work for bsz 64 right?
| elif workload_name == 'cifar': | ||
| return 128 | ||
| elif workload_name == 'lm': | ||
| return 8 |
There was a problem hiding this comment.
Should work for bsz 64 right?
4 participants
This is for the LM workload.